import torch

from problems.DTLZ import DTLZ
from problems.LSMOP import LSMOP


class DTLZ7(DTLZ):
    def __init__(self, var_dim: int, obj_dim: int, max_fun_eval, kwargs: dict):
        super().__init__(var_dim, obj_dim, max_fun_eval, kwargs)

    def eval_value(self, dec):
        # 初始化PopObj为零矩阵
        PopObj = torch.zeros(dec.size(0), self.obj_dim, dtype=torch.double)
        # 计算辅助函数g
        g = 1 + 9 * torch.mean(dec[:, self.obj_dim - 1:], dim=1)
        # 设置PopObj的前obj.M-1个目标
        PopObj[:, :self.obj_dim - 1] = dec[:, :self.obj_dim - 1]
        # 设置PopObj的第obj.M个目标
        # 首先，计算分母中的重复部分
        repeated_g = g.unsqueeze(1).repeat(1, self.obj_dim - 1)
        # 然后，计算最后一个目标的值
        PopObj[:, self.obj_dim - 1] = (1 + g) * (self.obj_dim - torch.sum(
            PopObj[:, :self.obj_dim - 1].clone() / (1 + repeated_g) * (
                    1 + torch.sin(3 * torch.pi * PopObj[:, :self.obj_dim - 1].clone())),
            dim=1))
        return PopObj

    def get_optimal_solutions(self, size):
        return LSMOP.get_optimal_solutions2(self.obj_dim, size)
